import swanlab
num_epochs = 20
lr = 1e-4
batch_size = 8
num_classes = 2
device = "cuda"
swanlab.init(
    # 设置实验名
    experiment_name="ResNet50",
    # 设置实验介绍
    description="Train ResNet50 for cat and dog classification.",
    # 记录超参数
    config={
        "model": "resnet50",
        "optim": "Adam",
        "lr": lr,
        "batch_size": batch_size,
        "num_epochs": num_epochs,
        "num_class": num_classes,
        "device": device,
    }
)
"--------------------------------数据载入部分--------------------------------"
import get_data
from torch.utils.data import DataLoader
train_dataset = get_data.DatasetLoader(get_data.ms_train_dataset)
train_loader = (DataLoader(train_dataset, batch_size=batch_size,shuffle=True))
"--------------------------------数据载入部分--------------------------------"
"--------------------------------model处理部分--------------------------------"
import torch
import torchvision
from torchvision.models import ResNet50_Weights
# 加载预训练的ResNet50模型
model = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
# 将全连接层的输出维度替换为num_classes
in_features = model.fc.in_features
model.fc = torch.nn.Linear(in_features, num_classes)
model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
"--------------------------------model处理部分--------------------------------"
for iter, (inputs, labels) in enumerate(train_loader):
    inputs, labels = inputs.to(device), labels.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'.format(num_epochs, num_epochs, iter + 1, len(train_loader), loss.item()))
    swanlab.log({"train_loss": loss.item()})
